#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Annotated
import re
from fastapi import APIRouter, Depends, Path, Query, Request, UploadFile

from app.schema.collection_schema import (
    CreateCollectionParam,
    UpdateCollectionParam,
    GetCollectionList, GetCollectionDetail
)
from app.schema.embedding_schema import CreateEmbeddingParam
from app.schema.text_block_schema import CreateTextBlockParam
from app.service.embedding_service import embedding_service
from app.service.collection_service import collection_service
from app.service.text_block_service import text_block_service
from common.pagination import DependsPagination, paging_data
from common.response.response_schema import ResponseModel, response_base
from common.security.jwt import DependsJwtAuth
from common.security.permission import RequestPermission
from core.conf import settings
from database.db_mysql import CurrentSession
from utils.sapper_server import send_async_rag_request
from utils.serializers import select_as_dict

router = APIRouter()


@router.get('/all', summary='查看知识库下所有集合', dependencies=[DependsJwtAuth])
async def get_all_collections(
        request: Request,
        knowledge_base_uuid: Annotated[str | None, Query()] = None,
) -> ResponseModel:
    # 获取所有会话
    agents = await collection_service.get_all(knowledge_base_uuid=knowledge_base_uuid)

    # 转换为响应所需的格式
    data = [GetCollectionList(**select_as_dict(agent)) for agent in agents]
    return response_base.success(data=data)


@router.post('', summary='添加知识集合', dependencies=[DependsJwtAuth])
async def add_collection(request: Request, obj: CreateCollectionParam) -> ResponseModel:
    collection = await collection_service.add(obj=obj)
    file = UploadFile(
        filename=obj.file_url.split("/")[-1],  # 提取文件名
        file=open(obj.file_url, "rb"),  # 以二进制模式打开文件
    )
    # 准备请求 URL 和 headers
    url = f'{settings.SAPPER_SERVER_URL}sapperrag/read'
    headers = {
        "Accept": "application/json"
    }

    # 发送异步请求
    embed_result = await send_async_rag_request(url, headers, file)
    # 插入 embeddings，使用单独事务减少锁等待超时问题
    for embed in embed_result:
        text_block_data = await text_block_service.add(
            obj=CreateTextBlockParam(
                collection_uuid=collection.uuid,
                content=embed.text
            )
        )
        await embedding_service.add(
            obj=CreateEmbeddingParam(
                vector=embed.text_embedding.tolist(),
                text_block_uuid=text_block_data.uuid
            )
        )

    collection_data = GetCollectionList(**select_as_dict(collection))
    return response_base.success(data=collection_data)


@router.put('/{collection_uuid}', summary='更新知识集合信息', dependencies=[DependsJwtAuth])
async def update_collection(request: Request, collection_uuid: Annotated[str, Path(...)], obj: UpdateCollectionParam) -> ResponseModel:
    count = await collection_service.update(request=request, collection_uuid=collection_uuid, obj=obj)
    if count > 0:
        return response_base.success()
    return response_base.fail()


@router.put('/{collection_uuid}/status', summary='修改知识集合状态', dependencies=[DependsJwtAuth])
async def update_collection_status(request: Request, collection_uuid: Annotated[str, Path(...)], status: Annotated[bool, Query()]) -> ResponseModel:
    count = await collection_service.update_status(request=request, collection_uuid=collection_uuid, status=status)
    if count > 0:
        return response_base.success()
    return response_base.fail()


@router.get('/{collection_uuid}', summary='查看知识集合信息', dependencies=[DependsJwtAuth])
async def get_collection_info(request: Request, collection_uuid: Annotated[str, Path(...)]) -> ResponseModel:
    current_collection = await collection_service.get_with_relation(request=request, collection_uuid=collection_uuid)
    data = GetCollectionDetail(**select_as_dict(current_collection))
    return response_base.success(data=data)


@router.get(
    '',
    summary='分页获取所有知识集合',
    dependencies=[
        DependsJwtAuth,
        DependsPagination,
    ],
)
async def get_pagination_collections(
    request: Request,
    db: CurrentSession,
    name: Annotated[str | None, Query()] = None,
    description: Annotated[str | None, Query()] = None,
    status: Annotated[int | None, Query()] = None,
):
    collection_select = await collection_service.get_select(user_uuid=request.user.uuid, name=name, description=description, status=status)
    page_data = await paging_data(db, collection_select, GetCollectionList)
    return response_base.success(data=page_data)


@router.delete(
    path='/{collection_uuid}',
    summary='删除知识集合',
    description='删除后知识集合将从数据库中删除',
    dependencies=[Depends(RequestPermission('sys:collection:del'))],
)
async def delete_collection(request: Request, collection_uuid: Annotated[str, Path(...)]) -> ResponseModel:
    count = await collection_service.delete(request=request, collection_uuid=collection_uuid)  # no need for request here
    if count > 0:
        return response_base.success()
    return response_base.fail()
